package fr.utbm.gi.vi51.slizers.framework.learning;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.TreeMap;
import java.util.Map.Entry;

/**
 * Implementation of a store for Q-Learning datas.
 * @author Tarik Belabbas
 * 
 */
public class QValueStore< StateType extends QState, ActionType extends QAction > implements Serializable
{
	/**
	 * 
	 */
	private static final long	serialVersionUID	= -6404981502490135317L;
	/**
	 * The tree map storing states and their associated actions/Q-values pairs.
	 */
	private final Map< StateType, Map< ActionType, Float > >	qValues	=
		new TreeMap< StateType, Map< ActionType, Float >>( new QComparator() );

	/**
	 * Parameterized constructor.
	 * 
	 * @param _problem the problem to use to construct the store
	 */
	public QValueStore( QProblem<StateType,ActionType> _problem )
	{
		for ( StateType currentState : _problem.getAvailableStates() )
		{
			Map< ActionType, Float > map = new TreeMap< ActionType, Float >(
					new QComparator() );
			this.qValues.put( currentState, map );

			for ( ActionType currentAction : _problem.getAvailableActionsFor( currentState ) )
			{
				map.put( currentAction, 0.0f );
			}
		}
	}

	/**
	 * Parameterized constructor.
	 * 
	 * @param _problem the problem to use to construct the store
	 */
	public QValueStore( QProblem< StateType, ActionType > _problem,
			QValueStore< StateType, ActionType > _qvaluestore )
	{
		for ( StateType currentState : _problem.getAvailableStates() )
		{
			Map< ActionType, Float > map = new TreeMap< ActionType, Float >(
					new QComparator() );
			this.qValues.put( currentState, map );
			
			Map< ActionType, Float > readMap = _qvaluestore.qValues.get( currentState );

			for ( ActionType currentAction : _problem.getAvailableActionsFor( currentState ) )
			{
				if ( readMap.containsKey( currentAction ) )
				{
					map.put( currentAction, readMap.get( currentAction ) );
				}
				else
				{
					System.out.println( "Invalid read QValueStore, not all current problem " +
							"actions were serialized (possible changes in available " +
							"actions" );
					map.put( currentAction, 0.0f );
				}
			}
		}
	}

	/**
	 * Returns the best action to take from a given state.
	 * 
	 * @param _state the state from which the best action must be determined
	 * @param _random the random object used in case of multiple best actions
	 * @return the best action
	 */
	public ActionType getBestAction( StateType _state, Random _random )
	{
		Map< ActionType, Float > map = this.qValues.get( _state );
		
		assert ( map != null );
		
		// Create a list of the best actions (if several)
		List< ActionType > bestActions = new ArrayList< ActionType >();
		float bestScore = Float.NEGATIVE_INFINITY;
		for ( Entry< ActionType, Float > currentEntry : map.entrySet() )
		{
			if ( currentEntry.getValue() > bestScore )
			{
				bestActions.clear();
				bestActions.add( currentEntry.getKey() );
				bestScore = currentEntry.getValue();
			}
			else if ( currentEntry.getValue() == bestScore )
			{
				bestActions.add( currentEntry.getKey() );
			}
		}
		
		assert ( !bestActions.isEmpty() );
		
		// If there are several best actions, pick up one randomly
		return bestActions.get( _random.nextInt( bestActions.size() ) );
	}
	
	public Map<Float, ActionType> getBestActions( StateType _state)
	{		
		Map< ActionType, Float > map = this.qValues.get( _state );
	
		assert ( map != null );
		Map< Float, ActionType > newMap = new TreeMap< Float, ActionType >();   
		for (ActionType key : map.keySet())   
		    newMap.put(map.get(key), key);
		return newMap;
	}

	/**
	 * Returns the current Q-value associated to a pair < State, Action >.
	 * 
	 * @param _state the state of the pair
	 * @param _action the action of the pair
	 * @return the current Q-value of the pair
	 */
	public float getQValue( StateType _state, ActionType _action )
	{
		Map< ActionType, Float > map = this.qValues.get( _state );
		
		assert ( map != null );
		
		Float qValue = map.get( _action );
		if ( qValue == null )
		{
			return 0.0f;
		}
		
		return qValue.floatValue();
	}
	
	/**
	 * Stores a Q-value for a given pair < State, Action >.
	 * 
	 * @param _state the state of the pair
	 * @param _action the action of the pair
	 * @param _qValue the new Q-value to store
	 */
	public void storeQValue( StateType _state, ActionType _action, float _qValue )
	{
		Map< ActionType, Float > map = this.qValues.get( _state );
		
		assert ( map != null );

		map.put( _action, _qValue );
	}
}
